import sys
from PyQt5.QtWidgets import (
    QApplication,
    QMainWindow,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QPushButton,
    QLineEdit,
    QFrame,
    QGridLayout,
    QComboBox,
    QLabel,
    QSlider,
    QMessageBox,
    QProgressBar,
    QRadioButton,
    QDialog,
)
import subprocess
from PyQt5.QtWidgets import QPlainTextEdit
from PyQt5.QtGui import QFontMetrics
from PyQt5.QtCore import Qt
from PyQt5 import QtWidgets, QtGui
from PyQt5.QtCore import QThread, pyqtSignal
from api import *
from PyQt5.QtGui import QPixmap, QImage
from PIL import Image
import os
from prompt_optimize import PromptOptimizeMachine, load_config
from wr import *
from page2 import *
from wr_lsb import *
from wr_dct import *
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding


class ImageGenerationThread(QThread):
    finished = pyqtSignal(str)
    error = pyqtSignal(str)

    def __init__(self, payload, pic_name):
        super().__init__()
        self.payload = payload
        self.pic_name = pic_name

    def run(self):
        try:
            get_pic_from_stable_diffusion(self.payload, self.pic_name)
            self.finished.emit(self.pic_name)
        except Exception as e:
            self.error.emit(str(e))


class ImageEmbeddingThread(QThread):
    finished = pyqtSignal(str, int)
    error = pyqtSignal(str)

    def __init__(self, image_path, data, fw):
        super().__init__()
        self.image_path = image_path
        self.data = data
        self.fw = fw

    def run(self):
        try:
            length = actor(self.image_path, self.data, self.fw)
            self.finished.emit("ans.png", length)
        except Exception as e:
            self.error.emit(str(e))


class Page1(QMainWindow):
    # get_pic = 0
    # image_path = 'example.png'
    def __init__(self, parent=None):
        super(Page1, self).__init__(parent)
        self.setWindowTitle("stable diffusion水印系统")
        self.setGeometry(100, 100, 1000, 800)
        ##
        # 创建中心部件
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        # 创建2x2网格布局
        grid_layout = QGridLayout(central_widget)
        grid_layout.setContentsMargins(10, 10, 10, 10)  # 设置边距为10
        grid_layout.setSpacing(10)  # 设置控件之间的间距

        # 创建四个模块并添加到网格布局中
        module1 = self.create_module_1()
        module2 = self.create_module_2()
        module3 = self.create_module_3()
        module4 = self.create_module_4()

        # 设置行和列的拉伸比例
        grid_layout.setRowStretch(0, 1)  # 第一行总宽度拉伸比例为6
        grid_layout.setRowStretch(0, 1)  # 第二行总宽度拉伸比例为6

        # 添加模块到布局中
        grid_layout.addWidget(module1, 0, 0, 1, 1)  # 第一行占两列，宽度比例3:2
        grid_layout.addWidget(module2, 0, 1, 1, 1)  # 第一行占两列，宽度比例3:2
        grid_layout.addWidget(module3, 1, 0, 1, 1)  # 第二行占一列，宽度比例1:1
        grid_layout.addWidget(module4, 1, 1, 1, 1)  # 第二行占一列，宽度比例1:1

        self.setStyleSheet(
            """  
                    QWidget {  
                background-color: #f0f0f0;  
            }  

            QLabel {  
                border: 2px solid #555;  
                border-radius: 5px;  
                padding: 10px;  
                font-size: 12pt;  
                font-family: "宋体", sans-serif; /* 设置统一的字体 */  
                color: #000; /* 设置统一的文本颜色 */  
                background-color: #fff;  
            }  
           QPushButton {  
                border: none;  
                border-radius: 5px;  
                padding: 10px 20px;  
                font-size: 12pt;  
                font-family: "Arial", sans-serif;  
                color: #fff;  
                background-color: #5cb85c; /* 更深的初始背景色 */    
            }  

            QPushButton:hover {  
                background-color: #4cae4c; /* 鼠标悬停时背景色稍浅 */  
            }  

            QPushButton:pressed {  
                background-color: #449d44; /* 按下时背景色更深 */  
                padding-top: 11px; /* 稍微调整内边距以模拟按下效果 */  
                padding-bottom: 9px;  
            }
           """
        )

        self.progress_bar = QProgressBar(self)
        self.statusBar().addWidget(self.progress_bar)
        self.progress_bar.setVisible(False)

    def create_module_1(self):
        # 第一个模块的布局和内容
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 创建水平布局用于放置文本输入框和下拉框
        layout_input = QHBoxLayout()

        # 创建第一个文本输入框并设置占位提示文字
        self.text_edit1 = QPlainTextEdit()
        self.text_edit1.setPlaceholderText("输入正向提示词（选项不为 None 时不支持中文提示词）")
        self.text_edit1.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")

        font_metrics = QFontMetrics(self.text_edit1.document().defaultFont())
        line_height = font_metrics.lineSpacing()

        visible_lines = 8  # 设置可见的行数
        desired_height = line_height * visible_lines
        self.text_edit1.setFixedHeight(desired_height)

        layout_input.addWidget(self.text_edit1)

        # 创建下拉框，并添加选项
        self.combo_box_prompt = QComboBox()
        self.combo_box_prompt.addItems(["None","翻译", "建筑", "肖像", "卡通", "游戏"])
        self.combo_box_prompt.setFixedHeight(desired_height)
        self.combo_box_prompt.setFixedWidth(70)  # 设置下拉框宽度
        self.combo_box_prompt.setStyleSheet("QComboBox { border: 2px solid black; }")

        layout_input.addWidget(self.combo_box_prompt)

        layout.addLayout(layout_input)

        # 创建第二个文本输入框并设置占位提示文字
        self.text_edit2 = QPlainTextEdit()
        self.text_edit2.setPlaceholderText("输入反向提示词")
        self.text_edit2.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")

        self.text_edit2.setFixedHeight(desired_height)
        layout.addWidget(self.text_edit2)
        # 创建水平布局用于放置滑块和标签
        layout1 = QHBoxLayout()
        layout1.setSpacing(5)  # 设置布局间距
        layout1.setAlignment(Qt.AlignCenter)  # 设置布局对齐方式

        # 创建滑块
        self.width_pic = QSlider(Qt.Horizontal)
        self.width_pic.setMinimum(0)  # 设置最小值
        self.width_pic.setMaximum(512)  # 设置最大值
        self.width_pic.setValue(256)  # 设置默认值（中心位置）
        self.width_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置
        self.width_pic.setTickInterval(10)  # 设置刻度间隔
        self.width_pic.setStyleSheet(
            "QSlider::groove:horizontal { border: 2px solid black; height: 10px; margin: 0; border-radius: 5px; } "  # 设置滑轨样式
            "QSlider::handle:horizontal { background: black; border: 2px solid black; width: 15px; margin: -7px 0; border-radius: 8px; }"
        )  # 设置滑块样式

        # 创建标签用于显示滑块的值
        self.label_width = QLabel("256")
        self.label_width.setAlignment(Qt.AlignCenter)  # 设置标签文本对齐方式

        # 创建标签用于显示'width : '文本
        lable_tmp = QLabel("width  : ")
        lable_tmp.setAlignment(Qt.AlignRight)  # 设置标签文本对齐方式，使其靠右显示

        # 将标签和滑块添加到水平布局中
        layout1.addWidget(lable_tmp)
        layout1.addWidget(self.width_pic)
        layout1.addWidget(self.label_width)

        # 将水平布局添加到垂直布局中（假设layout是已有的垂直布局）
        layout.addLayout(layout1)
        # 创建标签用于显示'width : '文本

        ##
        # 创建标签用于显示滑块的值
        # self.label_width = QLabel('256')
        # lable_tmp = QLabel('width  : ')

        # 将滑块的值改变信号与槽函数连接
        self.width_pic.valueChanged.connect(self.on_value_changed_width)
        layout1.addWidget(lable_tmp)
        layout1.addWidget(self.width_pic)
        layout1.addWidget(self.label_width)

        layout2 = QHBoxLayout()
        ##
        # 创建滑块
        self.height_pic = QSlider(Qt.Horizontal)
        self.height_pic.setMinimum(0)  # 设置最小值
        self.height_pic.setMaximum(512)  # 设置最大值
        self.height_pic.setValue(256)  # 设置默认值（中心位置）
        self.height_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置
        self.height_pic.setTickInterval(10)  # 设置刻度间隔

        # 设置滑轨和滑块样式
        self.height_pic.setStyleSheet(
            """  
            QSlider::groove:horizontal {  
                border: 2px solid black;  
                height: 10px;  
                margin: 0;  
                border-radius: 5px;  
            }  
            QSlider::handle:horizontal {  
                background: black;  
                border: 2px solid black;  
                width: 15px;  
                margin: -7px 0;  
                border-radius: 8px;  
            }  
        """
        )

        # 创建标签用于显示滑块的值
        self.label_height = QLabel("256")
        self.label_height.setAlignment(Qt.AlignCenter)  # 设置标签文本对齐方式

        # 创建标签用于显示'height : '文本
        lable_tmp = QLabel("height : ")
        lable_tmp.setAlignment(Qt.AlignRight)  # 设置标签文本对齐方式，使其靠右显示

        # 假设layout1是一个已经存在���水平布局
        layout2.addWidget(lable_tmp)
        layout2.addWidget(self.height_pic)
        layout2.addWidget(self.label_height)

        # 假设layout是一个已经存在的垂直布局
        # 将水平布局layout1添加到垂直布局layout中
        layout.addLayout(layout2)
        # 将滑块的值改变信号与槽函数连接
        self.height_pic.valueChanged.connect(self.on_value_changed_height)
        # 注意：下面这行代码不应该在这里，因为它会再次向layout1添加控件，这是重复的
        # layout1.addWidget(self.label_height)
        layout3 = QHBoxLayout()

        # 如果需要更新标签的值，可以在on_value_changed_height槽函数中实现
        def on_value_changed_height(self, value):
            self.label_height.setText(str(value))
            # 创建滑块

        self.step_pic = QSlider(Qt.Horizontal)
        self.step_pic.setMinimum(1)  # 设置最小值
        self.step_pic.setMaximum(40)  # 设置最大值
        self.step_pic.setValue(20)  # 设置默认值（中心位置）
        self.step_pic.setTickPosition(QSlider.TicksBelow)  # 设���刻度位置
        self.step_pic.setTickInterval(2)  # 设置刻度间隔

        # 设置滑轨和滑块样式
        self.step_pic.setStyleSheet(
            """  
            QSlider::groove:horizontal {  
                border: 2px solid black;  
                height: 10px;  
                margin: 0;  
                border-radius: 5px;  
            }  
            QSlider::handle:horizontal {  
                background: black;  
                border: 2px solid black;  
                width: 15px;  
                margin: -7px 0;  
                border-radius: 8px;  
            }  
        """
        )

        # # 创建标签用于显示滑块的值
        # self.label_step = QLabel('10')
        # self.label_step.setAlignment(Qt.AlignCenter)  # 设置标签文本对齐方式
        # # 创建标签用于显示'step : '文本
        # lable_tmp = QLabel('step : ')
        # lable_tmp.setAlignment(Qt.AlignRight)  # 设置标签文本对齐方式，使其靠右显示
        # 创建标签用于显示滑块的值
        self.label_step = QLabel("20")
        lable_tmp = QLabel("step : ")
        # 创建一个垂直布局来包裹标签、滑块和显示值的标签
        # layout3.addWidget(lable_tmp)
        # layout3.addWidget(self.step_pic)
        # layout3.addWidget(self.label_step)
        # 将滑块的值改变信号与槽函数连接
        self.step_pic.valueChanged.connect(self.on_value_changed_step)
        layout3.addWidget(lable_tmp)
        layout3.addWidget(self.step_pic)
        layout3.addWidget(self.label_step)

        layout.addLayout(layout1)
        layout.addLayout(layout2)
        layout.addLayout(layout3)
        widget.setLayout(layout)
        return widget

    def create_module_2(self):
        # 第二个模块的布局和内容
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 创建第一个文本输入框并设置占位提示文字
        self.text_edit3 = QPlainTextEdit()
        self.text_edit3.setPlaceholderText("输入嵌入数据")
        self.text_edit3.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")

        font_metrics = QFontMetrics(self.text_edit3.document().defaultFont())
        line_height = font_metrics.lineSpacing()

        visible_lines = 8  # 设置可见的行数
        desired_height = line_height * visible_lines
        self.text_edit3.setFixedHeight(desired_height)

        layout.addWidget(self.text_edit3)  # 将文本输入框添加到布局中

        # 创建下拉框，并添加选项
        lay1 = QHBoxLayout()
        label = QLabel("差异范围：")
        label.setFixedWidth(150)  # 设置固定宽度为120像素
        lay1.addWidget(label)
        self.combo_box = QComboBox()
        self.combo_box.addItems(["0-7", "8-15", "16-31", "32-63", "64-127", "128-255"])
        lay1.addWidget(self.combo_box)  # 将下拉框添加到布局中
        layout.addLayout(lay1)

        # 设置下拉框的大小，这里可以根据实际情况调整
        self.combo_box.setSizeAdjustPolicy(QComboBox.AdjustToContents)
        font = self.combo_box.font()
        font.setPointSize(15)  # 设置字体大小为15点
        self.combo_box.setFont(font)
        lay1.addWidget(self.combo_box)  # 将下拉框添加到布局中
        layout.addLayout(lay1)
        # 创建下拉框，并添加选项
        lay2 = QHBoxLayout()
        label = QLabel("checkpoint：")
        lay2.addWidget(label)
        self.combo_box1 = QComboBox()
        self.combo_box1.addItems(
            [
                "3Guofeng3_v34.safetensors",
                "blazingrealdrive_V02h.safetensors",
                "v2-1_768-ema-pruned.ckpt",
            ]
        )
        # 设置下拉框选项的字体大小
        font = self.combo_box1.font()
        font.setPointSize(12)  # 设置字体大小为12点，你可以根据需要调整这个值
        self.combo_box1.setFont(font)
        lay2.addWidget(self.combo_box1)  # 将下拉框添加到布局中
        layout.addLayout(lay2)

        layout3 = QHBoxLayout()

        layout2 = QHBoxLayout()
        label = QLabel("请输入保存文件名：")
        layout2.addWidget(label)
        self.pic_name = QLineEdit()
        self.pic_name.setStyleSheet("QLineEdit { border: 2px solid black; }")
        layout2.addWidget(self.pic_name)
        layout.addLayout(layout2)

        # 创建滑块
        self.CFG_pic = QSlider(Qt.Horizontal)
        self.CFG_pic.setMinimum(1)  # 设置最小值
        self.CFG_pic.setMaximum(10)  # 设置最大值
        self.CFG_pic.setValue(5)  # 设置默认值（中心位置）
        self.CFG_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置
        self.CFG_pic.setTickInterval(1)  # 设置刻度间隔
        # 设置滑轨和滑块样式
        self.CFG_pic.setStyleSheet(
            """  
            QSlider::groove:horizontal {  
                border: 2px solid black;  
                height: 10px;  
                margin: 0;  
                border-radius: 5px;  
            }  
            QSlider::handle:horizontal {  
                background: black;  
                border: 2px solid black;  
                width: 15px;  
                margin: -7px 0;  
                border-radius: 8px;  
            }  
        """
        )
        # 创建标签用于显示滑块的值
        self.label_CFG = QLabel("5")
        lable_tmp = QLabel("CFG Scale : ")
        # 将滑块的值改变信号与槽函数连接
        self.CFG_pic.valueChanged.connect(self.on_value_changed_CFG)
        layout3.addWidget(lable_tmp)
        layout3.addWidget(self.CFG_pic)
        layout3.addWidget(self.label_CFG)
        layout.addLayout(layout3)

        # 创建两个按钮，并排对称排布
        button1 = QPushButton("生成图片")
        button2 = QPushButton("嵌入水印")
        button3 = QPushButton("水印提取")
        button4 = QPushButton("生成密钥")
        button1.clicked.connect(self.on_button_clicked1)
        button2.clicked.connect(self.on_button_clicked2)
        button3.clicked.connect(self.on_button_clicked3)
        button4.clicked.connect(self.on_button_clicked4)
        # central_widget = QtWidgets.QWidget()
        # self.page2 = Page2(central_widget)

        # 为了使两个按钮对称排布，我们可以将它们放入一个水平布局中，然后将这个布局添加到主布局中
        sub_layout = QHBoxLayout()
        sub_layout.addWidget(button1)
        sub_layout.addWidget(button2)
        sub_layout.addWidget(button3)
        sub_layout.addWidget(button4)
        sub_layout.setAlignment(Qt.AlignCenter)  # 设置按钮居中对齐

        layout.addLayout(sub_layout)  # 将按钮布局添加到主布局中

        widget.setLayout(layout)
        return widget

    def create_module_3(self):
        # 第三个模块的布局和内容
        image_path = "C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\SD\\LN.jpg"
        widget = QWidget()
        layout = QHBoxLayout(widget)

        # 创建一个QLabel对象用于显示图片
        self.image_label_init = QLabel()

        # 使用QPixmap加载图片
        pixmap = QtGui.QPixmap(image_path)

        # 如果图片大小大于标签当前大小，则调整标签大小以适应图片
        self.image_label_init.setPixmap(
            pixmap.scaled(
                self.image_label_init.size(), aspectRatioMode=Qt.KeepAspectRatio
            )
        )

        # 将QLabel添加到布局中
        layout.addWidget(self.image_label_init)

        # 设置布局到widget
        widget.setLayout(layout)
        return widget

    def create_module_4(self):
        # 第三个模块的布局和内容
        image_path = "C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\SD\\LN1.png"
        widget = QWidget()
        layout = QHBoxLayout(widget)

        # 创建一个QLabel对象用于显示图片
        self.image_label_ans = QLabel()

        # 使用QPixmap加载图片
        pixmap = QtGui.QPixmap(image_path)

        # 如果图片大小大于标签当前大小，则调整标签大小以适应图片
        self.image_label_ans.setPixmap(
            pixmap.scaled(
                self.image_label_ans.size(), aspectRatioMode=Qt.KeepAspectRatio
            )
        )

        # 将QLabel添加到布局中
        layout.addWidget(self.image_label_ans)

        # 设置布局到widget
        widget.setLayout(layout)
        return widget

    def on_value_changed_width(self, value):
        # 更新标签显示的值
        self.label_width.setText(str(value))

    def on_value_changed_height(self, value):
        # 更新标签显示的值
        self.label_height.setText(str(value))

    def on_value_changed_step(self, value):
        # 更新标签显示的值
        self.label_step.setText(str(value))

    def on_value_changed_CFG(self, value):
        # 更新标签显示的值
        self.label_CFG.setText(str(value))

    def on_button_clicked3(self):
        page22 = Page2(self)
        page22.show()

    def on_button_clicked4(self):
        # 生成RSA密钥对
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
        )

        # 获取公钥
        public_key = private_key.public_key()

        # 序列化私钥和公钥，保存为PEM格式文件
        with open("private_key.pem", "wb") as private_key_file:
            private_key_file.write(private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            ))

        with open("public_key.pem", "wb") as public_key_file:
            public_key_file.write(public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            ))

    def on_button_clicked1(self):
        width = self.width_pic.value()
        heigth = self.height_pic.value()
        step = self.step_pic.value()
        cfg = self.CFG_pic.value()
        negative_prompt = self.text_edit2.toPlainText()
        prompt = self.text_edit1.toPlainText()
        checkpoint = self.combo_box1.currentText()
        if width < 8 or heigth < 8:
            QMessageBox.warning(self, "警告", "长或宽不能小于8个像素")
            return
        if not prompt and not negative_prompt:
            # 如果为空，显示提示框
            QMessageBox.warning(self, "警告", "正向/反向提示词不能同时为空")
            # 退出函数，不再执行后续代码
            return

        # TODO: 添加对提示词的优化

        # 加载 API 密钥
        config = load_config("config.conf")
        os.environ["TENCENTCLOUD_SECRET_ID"] = config.get("SecretId")
        os.environ["TENCENTCLOUD_SECRET_KEY"] = config.get("SecretKey")
        
        # 优化提示词
        pom = PromptOptimizeMachine(prompt, negative_prompt)
        selected_value = self.combo_box_prompt.currentText()
        prompt, negative_prompt = pom.optimize(selected_value)
        
        if not negative_prompt:
            negative_prompt = ""

        payload = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "steps": step,
            "width": width,
            "height": heigth,
            "seed": -1,
            "CFG Scale": cfg,
            "stable diffusion checkpoint": checkpoint,  # 下拉框
        }
        
        print(payload)
        
        self.get_pic = 1
        # 更换图片
        # 使用Pillow打开图片
        self.image_path = self.pic_name.text()
        if len(self.image_path) == 0:
            self.image_path = "example.png"
        else:
            self.image_path = self.image_path + ".png"
            # print(payload)
        print(self.image_path)
        # get_pic_from_stable_diffusion(payload,self.image_path)
        # pil_image = Image.open(self.image_path)
        self.progress_bar.setRange(0, 0)
        self.progress_bar.setVisible(True)
        # self.progress_bar.setValue(0)
        self.thread = ImageGenerationThread(payload, self.image_path)
        self.thread.finished.connect(self.on_generation_finished)
        self.thread.error.connect(self.on_generation_error)
        self.thread.start()

    def on_button_clicked2(self):
        # 创建选择对话框
        dialog = QDialog(self)
        dialog.setWindowTitle("选择嵌入方式")
        dialog.setWindowFlags(dialog.windowFlags() & ~Qt.WindowContextHelpButtonHint)
        dialog.setFixedSize(450, 300)
        layout = QVBoxLayout(dialog)

        # label = QLabel("嵌入方式选择")
        # label.setAlignment(Qt.AlignCenter)  # 设置标签居中对齐
        # layout.addWidget(label)

        # 创建LSB和PVD、DCT的单选按钮
        self.lsb_radio = QRadioButton("LSB")
        self.pvd_radio = QRadioButton("PVD")
        self.dct_radio = QRadioButton("DCT")
        self.lsb_radio.setChecked(True)  # 默认选择LSB

        layout.addWidget(self.lsb_radio)
        layout.addWidget(self.pvd_radio)
        layout.addWidget(self.dct_radio)

        # 添加伸缩项以使按钮居中
        layout.addStretch()
        # 创建确认按钮
        confirm_button = QPushButton("确认")
        confirm_button.clicked.connect(dialog.accept)
        layout.addWidget(confirm_button)

        if dialog.exec_() == QDialog.Accepted:
            # 获取用户选择的嵌入方式
            if self.lsb_radio.isChecked():
                embedding_method = "LSB"
            elif self.pvd_radio.isChecked():
                embedding_method = "PVD"
            elif self.dct_radio.isChecked():
                embedding_method = "DCT"
            data = self.text_edit3.toPlainText()
            choose = self.combo_box.currentText()
            choose = choose.split("-")
            min_d = int(choose[0])
            max_d = int(choose[1])
            fw = (min_d, max_d)
            if len(data) == 0:
                QMessageBox.warning(self, "警告", "嵌入数据不可为空")
                return
            if not os.path.exists(self.image_path):
                QMessageBox.warning(self, "警告", "图片尚未生成")
                return
            # rsa修改
            # 加载公钥
            with open("public_key.pem", "rb") as public_key_file:
                public_key = serialization.load_pem_public_key(public_key_file.read())
            # RSA 加密
            encrypted_message = public_key.encrypt(
                data,
                padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
                )
            )
            # 根据选择的嵌入方式调用不同的处理函数
            if embedding_method == "LSB":
                # rsa修改
                l = actor_lsb(self.image_path, encrypted_message)  # 假设您有一个LSB的处理函数
                QMessageBox.warning(self, "提示", "加密长度为： " + str(l))
                self.image_path = "ans_lsb.png"
                pil_image = Image.open(self.image_path)

                try:
                    # 将Pillow图像转换为QImage
                    qimage = QImage(
                        pil_image.tobytes(),
                        pil_image.size[0],
                        pil_image.size[1],
                        QImage.Format.Format_RGB888,
                    )
                except Exception as e:
                    print("An error occurred:", e)
                # 将QImage转换为QPixmap
                try:
                    pixmap = QPixmap.fromImage(qimage)
                except Exception as e:
                    print("An error occurred:", e)
                # 在QLabel中显示灰度图片
                self.image_label_ans.setPixmap(pixmap)
            elif embedding_method == "DCT":
                try:
                    # watermark_size, synthesis = actor_dct(self.image_path, data)
                    # np.save("synthesis_data.npy", synthesis)
                    # rsa修改
                    watermark_size = actor_dct(self.image_path, encrypted_message)

                    # 提示水印嵌入后的��度或尺寸
                    QMessageBox.warning(
                        self, "提示", f"加密水印尺寸为： {watermark_size}"
                    )

                    # 更新嵌入后的图像路径
                    self.image_path = "ans_dct.png"

                    # 使用 Pillow 打开嵌入后的图片
                    pil_image = Image.open(self.image_path).convert(
                        "RGB"
                    )  # 确保是 RGB 格式

                    # 将 Pillow 图像转换为 QImage
                    try:
                        qimage = QImage(
                            pil_image.tobytes(),
                            pil_image.size[0],
                            pil_image.size[1],
                            pil_image.size[0] * 3,
                            QImage.Format.Format_RGB888,
                        )
                    except Exception as e:
                        print("QImage 转换错误:", e)

                    # 将 QImage 转换为 QPixmap
                    try:
                        pixmap = QPixmap.fromImage(qimage)
                    except Exception as e:
                        print("QPixmap 转换错误:", e)

                    # 在 QLabel 中显示嵌入后的图片
                    self.image_label_ans.setPixmap(pixmap)
                except Exception as e:
                    QMessageBox.warning(self, "错误", f"加密过程中发生错误：{str(e)}")

            else:
                self.progress_bar.setRange(0, 0)
                self.progress_bar.setVisible(True)
                # rsa修改
                l = actor(self.image_path, encrypted_message, fw)  # 假设您有一个PVD的处理函数
                QMessageBox.warning(self, "提示", "加密长度为： " + str(l))
                self.thread = ImageEmbeddingThread(self.image_path, data, fw)
                self.thread.finished.connect(self.on_embedding_finished)
                self.thread.error.connect(self.on_embedding_error)
                self.thread.start()
            # 使用Pillow打开图片

    def on_generation_finished(self, image_path):
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(100)
        self.progress_bar.setVisible(False)
        self.display_generated_image(image_path)
        QMessageBox.information(self, "成功", "图片生成成功！")

    def on_generation_error(self, error_message):
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setVisible(False)
        QMessageBox.critical(self, "错误", f"生成图片时出错: {error_message}")

    def display_generated_image(self, image_path):
        pil_image = Image.open(f"images/{image_path}")
        try:
            qimage = QImage(
                pil_image.tobytes(),
                pil_image.size[0],
                pil_image.size[1],
                QImage.Format_RGB888,
            )
            pixmap = QPixmap.fromImage(qimage)
            self.image_label_init.setPixmap(pixmap)
        except Exception as e:
            print("An error occurred:", e)

    def on_embedding_finished(self, image_path, length):
        # 恢复进度条为确定模式并隐藏
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(100)
        self.progress_bar.setVisible(False)

        # 显示嵌入后的图片
        self.image_path = image_path
        self.display_embedded_image(image_path)

    def on_embedding_error(self, error_message):
        # 恢复进度条为确定模式并隐藏
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setVisible(False)
        QMessageBox.critical(self, "错误", f"嵌入水印时出错: {error_message}")

    def display_embedded_image(self, image_path):
        pil_image = Image.open(image_path)
        try:
            qimage = QImage(
                pil_image.tobytes(),
                pil_image.size[0],
                pil_image.size[1],
                QImage.Format_RGB888,
            )
            pixmap = QPixmap.fromImage(qimage)
            self.image_label_ans.setPixmap(pixmap)
        except Exception as e:
            print("An error occurred:", e)

    def display_error(self, error_message):
        QMessageBox.critical(self, "错误", f"生成图片时出错: {error_message}")


if __name__ == "__main__":
    if os.name == 'nt':
        subprocess.run(["powershell", "-File", "SetKey.ps1"])
    app = QApplication(sys.argv)
    main_window = Page1()
    main_window.show()
    sys.exit(app.exec_())
